package sit.engreal.leetcode.problems100;

import org.junit.jupiter.api.Assertions;

/**
 * <a herf="https://leetcode.cn/problems/median-of-two-sorted-arrays/">4. 寻找两个正序数组的中位数</a>
 */
public class LeetCode4 {
    public static class Solution {

        private int findKth(int[] nums1, int[] nums2, int k) {
            int offset1 = 0;
            int offset2 = 0;
            for (; ; ) {
                if (nums1.length - offset1 > nums2.length - offset2) {
                    int[] a = nums1;
                    nums1 = nums2;
                    nums2 = a;

                    int n = offset1;
                    offset1 = offset2;
                    offset2 = n;
                }
                if (nums1.length - offset1 <= 0) {
                    return nums2[offset2 + k - 1];
                }
                if (k == 1) {
                    return Math.min(nums1[offset1], nums2[offset2]);
                }
                int i = offset1 + k / 2 - 1;
                int j = offset2 + k / 2 - 1;
                if (nums1.length <= i) {
                    offset2 = j + 1;
                } else if (nums1[i] < nums2[j]) {
                    offset1 = i + 1;
                } else {
                    offset2 = j + 1;
                }
                k -= (k / 2);
            }
        }

        public double findMedianSortedArrays(int[] nums1, int[] nums2) {
            final int left = (nums1.length + nums2.length + 1) / 2;
            final int right = (nums1.length + nums2.length + 2) / 2;
            final int a = findKth(nums1, nums2, left);
            if (left == right) {
                return a;
            } else {
                final int b = findKth(nums1, nums2, right);
                return (a + b) * 0.5;
            }
        }

        public static void main(String[] args) {
            final Solution solution = new Solution();

            Assertions.assertEquals(3.5,
                    solution.findMedianSortedArrays(
                            new int[]{4},
                            new int[]{1, 2, 3, 5, 6}
                    ), 0.00001);

            Assertions.assertEquals(2.0,
                    solution.findMedianSortedArrays(
                            new int[]{1, 2},
                            new int[]{2}
                    ), 0.00001);

            Assertions.assertEquals(2.5,
                    solution.findMedianSortedArrays(
                            new int[]{1, 2},
                            new int[]{3, 4}
                    ), 0.00001);
        }
    }
}
